import type { ItemType } from 'antd/es/menu/interface'
import { create } from 'zustand'
import React from 'react'

/** 菜单定义 */
export type MenuItem = ItemType & {
  /** 主键，亦用作菜单路由 */
  key?: string
  /** 标题 */
  label?: string
  /** 分组 */
  type?: 'group' | 'divider'
  /** 图标 */
  icon?: React.ReactNode
  /** 子菜单 */
  children?: MenuItem[]
  /**
   * 标签栏行为
   * - hidden: 隐藏
   */
  tabBar?: 'hidden'
}

type MenuStore = {
  /** 菜单定义 */
  menus: MenuItem[]
  /**
   * 菜单Map，提供一些方法用于快速查找菜单信息
   * - 在setMenus时进行创建
   */
  menuMap: MenuItemMap
  /**
   * 更新菜单
   * @param menus 菜单定义
   */
  setMenus: (menus: MenuItem[]) => void
}

/**
 * 菜单Map
 */
class MenuItemMap {
  /**
   * 菜单索引
   */
  private value: Map<string, { node: MenuItem; parent?: MenuItem }>

  /**
   * 初始化菜单Map
   * @param menus 菜单定义集合
   */
  constructor(menus: MenuItem[]) {
    this.value = new Map()
    this.buildMenuItemMap(menus)
  }

  /**
   * 根据主键查找对应的菜单
   * @param key 主键
   */
  public get(key: string): MenuItem | undefined {
    return this.value.get(key)?.node
  }

  /**
   * 根据主键查找对应的菜单父层
   * @param key 主键
   */
  public getParent(key: string): MenuItem | undefined {
    return this.value.get(key)?.parent
  }

  /**
   * 根据索引查找父级路径，从根节点到父级的完整路径
   * - 比如 /系统/产品管理/新增 -> /系统/产品管理
   * @param key 目标key
   */
  public getParentKeys(key: string) {
    const path: string[] = []
    let current = this.value.get(key)

    while (current && current.parent) {
      // 添加到路径开头
      const parentKey = current.parent.key!
      path.unshift(parentKey)
      current = this.value.get(parentKey)
    }
    return path
  }

  /**
   * 根据索引查找菜单路径
   * - 比如 /系统/产品管理/新增 -> /系统/产品管理/新增
   * @param key 目标key
   */
  public getMenuPaths(key: string) {
    const path: MenuItem[] = []
    let current = this.value.get(key)
    if (current) {
      path.push(current.node)
    }
    while (current && current.parent) {
      // 添加到路径开头
      const parentKey = current.parent.key!
      path.unshift(current.parent)
      current = this.value.get(parentKey)
    }
    return path
  }

  /**
   * 根据主键构建菜单索引，将所有菜单根据主键平铺
   * @param menus 菜单
   */
  private buildMenuItemMap(menus: MenuItem[]) {
    const index = new Map()

    function traverse(nodes: MenuItem[], parent?: MenuItem) {
      nodes.forEach((node) => {
        index.set(node.key, { node, parent })
        if (node.children) {
          traverse(node.children, node)
        }
      })
    }

    traverse(menus, undefined)
    this.value = index
  }
}

/**
 * 菜单状态库，存储菜单定义
 */
export const useMenuStore = create<MenuStore>((set) => ({
  menus: [],
  menuMap: new MenuItemMap([]),
  setMenus: (menus) => {
    const menuMap = new MenuItemMap(menus)
    set({ menus, menuMap })
  },
}))
